
#include "mat_mul_sub_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
using namespace matmul_tiling;
namespace optiling {
const uint32_t BLOCK_SIZE = 32;

int DoMmTiling(int M, int N, int K, int baseM, int baseN, int usedCoreNum, uint32_t sizeofdatatype,
               TCubeTiling *tiling) {
  TPosition leftPos = TPosition::GM;
  CubeFormat leftFormat = CubeFormat::ND;
  matmul_tiling::DataType leftDtype = (sizeofdatatype == 2) ? matmul_tiling::DataType::DT_FLOAT16 : matmul_tiling::DataType::DT_FLOAT;
  int transposeA = 0;

  TPosition rightPos = TPosition::GM;
  CubeFormat rightFormat = CubeFormat::ND;
  matmul_tiling::DataType rightDtype = (sizeofdatatype == 2) ? matmul_tiling::DataType::DT_FLOAT16 : matmul_tiling::DataType::DT_FLOAT;
  int transposeB = 0;

  TPosition resPos = TPosition::GM;
  CubeFormat resFormat = CubeFormat::ND;
  matmul_tiling::DataType resDtype = (sizeofdatatype == 2) ? matmul_tiling::DataType::DT_FLOAT16 : matmul_tiling::DataType::DT_FLOAT;

  TPosition biasPos = TPosition::GM;
  CubeFormat biasFormat = CubeFormat::ND;
  matmul_tiling::DataType biasDtype = (sizeofdatatype == 2) ? matmul_tiling::DataType::DT_FLOAT16 : matmul_tiling::DataType::DT_FLOAT;
  int isBias = 0;

  tiling->set_usedCoreNum(usedCoreNum);
  MultiCoreMatmulTiling tilingApi;
  tilingApi.SetDim(usedCoreNum);
  tilingApi.SetAType(leftPos, leftFormat, leftDtype, bool(transposeA));
  tilingApi.SetBType(rightPos, rightFormat, rightDtype, bool(transposeB));
  tilingApi.SetCType(resPos, resFormat, resDtype);
  tilingApi.SetBiasType(biasPos, biasFormat, biasDtype);
  tilingApi.SetOrgShape(M, N, K);
  tilingApi.SetShape(M, N, K);
  tilingApi.SetBias(bool(isBias));
  tilingApi.SetFixSplit(baseM, baseN, -1);
  tilingApi.SetBufferSpace(-1, -1, -1);
  tilingApi.SetSingleShape(M, N, K);

  int64_t res = tilingApi.GetTiling(*tiling);
  if (res == -1) {
    std::cout << "DoMmTiling fail" << std::endl;
    return res;
  }
  tiling->set_iterateOrder(1);

  uint32_t baseK = tiling->get_baseK();
  uint32_t bestL1Part1 = 160 * 1024;
  uint32_t bestL1Part2 = 128 * 1024;
  uint32_t stepKa = (bestL1Part1 >> 1) / (baseM * baseK * sizeofdatatype);
  uint32_t stepKb = (bestL1Part2 >> 1) / (baseN * baseK * sizeofdatatype);
  if (stepKa > stepKb) {
    stepKa = stepKa / stepKb * stepKb;
  } else if (stepKa < stepKb) {
    stepKb = stepKb / stepKa * stepKa;
  }
  tiling->set_stepKa(stepKa);
  tiling->set_stepKb(stepKb);
  tiling->set_stepM(1);
  tiling->set_stepN(1);
  tiling->set_depthA1(2 * stepKa);
  tiling->set_depthB1(2 * stepKb);
  return 0;
}

static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
  MatMulSubTilingData tiling;
    uint32_t sizeofdatatype;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto aicNum = ascendcPlatform.GetCoreNumAic();
    auto aivNum = ascendcPlatform.GetCoreNumAiv();
    context->SetBlockDim(aicNum);
    auto dataType = context->GetInputTensor(0)->GetDataType();
    if(dataType == ge::DT_UINT8) {
        sizeofdatatype = 1;
    }else if(dataType == ge::DT_FLOAT16) {
        sizeofdatatype = 2;
    }else{
        sizeofdatatype = 4;
    }
    uint32_t alignNum = BLOCK_SIZE / sizeofdatatype;
    int M = context->GetInputShape(0)->GetStorageShape().GetDim(0);
    int K = context->GetInputShape(0)->GetStorageShape().GetDim(1);
    int N = context->GetInputShape(1)->GetStorageShape().GetDim(1);
    int dimX3 = context->GetInputShape(2)->GetStorageShape().GetDimNum();
    int baseM = 128;
    int baseN = (N + 15) / 16 * 16;
    baseN = baseN > 256 ? 256 : baseN;
    int singleM = 128;
    int singleN = 128;

    DoMmTiling(baseM, N, K, baseM, baseN, aivNum, sizeofdatatype, &tiling.mmTiling);

    tiling.set_usedCoreNum(aicNum);
    tiling.set_alignNum(alignNum);
    tiling.set_M(M);
    tiling.set_N(N);
    tiling.set_K(K);
    tiling.set_dimX3(dimX3);
    tiling.set_singleM(singleM);
    tiling.set_singleN(singleN);

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    size_t userWorkspaceSize = M * N * sizeofdatatype;
    size_t systemWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = userWorkspaceSize + systemWorkspaceSize;
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class MatMulSub : public OpDef {
public:
    explicit MatMulSub(const char* name) : OpDef(name)
    {
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x3")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");
    }
};

OP_ADD(MatMulSub);
}
